import os
import sys
import pickle
from pathlib import Path

import pandas as pd
import matplotlib.pyplot as plt

BASE = Path(__file__).resolve().parents[1]
REPORT_DIR = BASE / 'reports'
REPORT_DIR.mkdir(exist_ok=True)

# Find latest pred_best_*.pkl
candidates = sorted(BASE.glob('pred_best_*.pkl'), key=lambda p: p.stat().st_mtime, reverse=True)
if not candidates:
    print('No pred_best_*.pkl found under', BASE)
    sys.exit(1)
pred_path = candidates[0]
print('Using prediction file:', pred_path)
run_id = pred_path.stem.replace('pred_best_', '')

# Locate artifacts dir for label.pkl
artifacts_dir = None
for exp_dir in (BASE / 'mlruns').iterdir():
    if not exp_dir.is_dir() or exp_dir.name == '.trash':
        continue
    candidate = exp_dir / run_id / 'artifacts'
    if candidate.exists():
        artifacts_dir = candidate
        break
if artifacts_dir is None:
    print('Artifacts dir not found for run:', run_id)
    sys.exit(2)

label_path = artifacts_dir / 'label.pkl'
if not label_path.exists():
    print('label.pkl not found at', label_path)
    sys.exit(3)

with open(pred_path, 'rb') as f:
    pred = pickle.load(f)
with open(label_path, 'rb') as f:
    label = pickle.load(f)

# Normalize to DataFrame with MultiIndex
if isinstance(pred, pd.Series):
    pred = pred.to_frame('score')
if isinstance(label, pd.Series):
    label = label.to_frame('label')
if 'label' not in label.columns:
    # try to find the label column
    if len(label.columns) == 1:
        label.columns = ['label']
    else:
        label = label.rename(columns={label.columns[0]: 'label'})

# Align indexes
df = pred.join(label, how='inner')
if df.empty:
    print('Joined prediction-label DataFrame is empty; index mismatch.')
    sys.exit(4)

# Create per-run report directory
run_report_dir = REPORT_DIR / run_id
run_report_dir.mkdir(parents=True, exist_ok=True)

# 1) Score histogram
plt.figure(figsize=(8, 5))
df['score'].hist(bins=50)
plt.title('Score Distribution')
plt.xlabel('score')
plt.ylabel('count')
plt.tight_layout()
plt.savefig(run_report_dir / 'score_hist.png')
plt.close()

# 2) Daily IC (Pearson)
try:
    ic_daily = df.groupby(level=0).apply(lambda x: x['score'].corr(x['label']))
except Exception:
    # fallback using numpy corrcoef
    def corr_np(x):
        import numpy as np
        a = x['score'].values
        b = x['label'].values
        if len(a) < 2:
            return float('nan')
        return float(np.corrcoef(a, b)[0, 1])
    ic_daily = df.groupby(level=0).apply(corr_np)

ic_mean = ic_daily.mean()
ic_std = ic_daily.std()
ic_ir = ic_mean / (ic_std if ic_std != 0 else float('nan'))

plt.figure(figsize=(10, 4))
ic_daily.plot()
plt.title(f'Daily IC (mean={ic_mean:.4f}, IR={ic_ir:.3f})')
plt.xlabel('date')
plt.ylabel('IC')
plt.tight_layout()
plt.savefig(run_report_dir / 'ic_daily.png')
plt.close()

# 3) Daily Rank IC
try:
    ric_daily = df.groupby(level=0).apply(lambda x: x['score'].rank().corr(x['label'].rank()))
    ric_mean = ric_daily.mean()
    ric_std = ric_daily.std()
    ric_ir = ric_mean / (ric_std if ric_std != 0 else float('nan'))
    plt.figure(figsize=(10, 4))
    ric_daily.plot(color='orange')
    plt.title(f'Daily Rank IC (mean={ric_mean:.4f}, IR={ric_ir:.3f})')
    plt.xlabel('date')
    plt.ylabel('Rank IC')
    plt.tight_layout()
    plt.savefig(run_report_dir / 'rank_ic_daily.png')
    plt.close()
except Exception as e:
    print('Rank IC computation failed:', repr(e))

# 4) Long-Short cumulative return (top 10% - bottom 10%)
ls_series = []
for date, x in df.groupby(level=0):
    x_sorted = x.sort_values('score', ascending=False)
    n = max(int(len(x_sorted) * 0.1), 1)
    top_mean = x_sorted.head(n)['label'].mean()
    bottom_mean = x_sorted.tail(n)['label'].mean()
    ls_series.append((date, top_mean - bottom_mean))
ls_df = pd.DataFrame(ls_series, columns=['date', 'ls_ret']).set_index('date').sort_index()
ls_cum = ls_df['ls_ret'].cumsum()
plt.figure(figsize=(10, 4))
ls_cum.plot(color='green')
plt.title('Long-Short Cumulative Return (Top 10% - Bottom 10%)')
plt.xlabel('date')
plt.ylabel('cum return')
plt.tight_layout()
plt.savefig(run_report_dir / 'long_short_cumret.png')
plt.close()

# 5) Top-50 cumulative return
top50_series = []
for date, x in df.groupby(level=0):
    x_sorted = x.sort_values('score', ascending=False)
    k = min(50, len(x_sorted))
    top50_mean = x_sorted.head(k)['label'].mean()
    top50_series.append((date, top50_mean))
top50_df = pd.DataFrame(top50_series, columns=['date', 'ret']).set_index('date').sort_index()
top50_cum = top50_df['ret'].cumsum()
plt.figure(figsize=(10, 4))
top50_cum.plot(color='purple')
plt.title('Top-50 Cumulative Return (equal-weight)')
plt.xlabel('date')
plt.ylabel('cum return')
plt.tight_layout()
plt.savefig(run_report_dir / 'top50_cumret.png')
plt.close()

# Summary CSV
summary = pd.DataFrame({
    'metric': ['IC.mean', 'IC.std', 'IC.IR', 'RankIC.mean', 'RankIC.std', 'RankIC.IR'],
    'value': [ic_mean, ic_std, ic_ir, ric_mean if 'ric_daily' in locals() else float('nan'), ric_std if 'ric_daily' in locals() else float('nan'), ric_ir if 'ric_daily' in locals() else float('nan')],
})
summary.to_csv(run_report_dir / 'summary_metrics.csv', index=False)

print('Saved visualizations to:', run_report_dir)